package org.lql.chaos.kafka.example.interceptor;

import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.header.Header;
import org.lql.chaos.kafka.example.context.PartitionStrategyContext;

import java.nio.charset.StandardCharsets;
import java.util.Map;

/**
 * @author: liangqinglong
 * @date: 2025-08-10 14:59
 * @description:
 **/
public class PartitionStrategyInterceptor implements ProducerInterceptor<String, Object> {

	@Override
	public ProducerRecord<String, Object> onSend(ProducerRecord<String, Object> record) {
		Header header = record.headers().lastHeader("partition-strategy");
		if (header != null) {
			String strategy = new String(header.value(), StandardCharsets.UTF_8);
			PartitionStrategyContext.setStrategy(strategy);
		} else {
			PartitionStrategyContext.setStrategy("DEFAULT");
		}
		return record;
	}

	@Override
	public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
		PartitionStrategyContext.clear();
	}

	@Override
	public void close() {
		PartitionStrategyContext.clear();
	}

	@Override
	public void configure(Map<String, ?> configs) {
	}
}
